package org.usmile.algorithms.learn.skiplist;

import java.util.Random;

public class SkipList<E extends Comparable<E>> {
    private static final int MAX_LEVEL = 16;

    private static class Node<E extends Comparable<E>> {
        E e;

        Node<E>[] next;

        @SuppressWarnings("all")
        public Node(E e) {
            this.e = e;

            this.next = (Node<E>[]) new Node[MAX_LEVEL];
        }
    }

    private int maxLevel;
    private final Random random = new Random();

    private final Node<E> dummyHead;

    public SkipList() {
        this.maxLevel = 1;
        this.dummyHead = new Node<>(null);
    }

    public boolean contains(E e) {
        return getNode(e) != null;
    }

    private Node<E> getNode(E e) {
        Node<E> prev = dummyHead;
        for (int i = maxLevel - 1; i >= 0; i--) {
            while (prev.next[i] != null && prev.next[i].e.compareTo(e) < 0) {
                prev = prev.next[i];
            }
        }

        if (prev.next[0] != null && prev.next[0].e.compareTo(e) == 0) {
            return prev.next[0];
        }

        return null;
    }

    @SuppressWarnings("all")
    public void add(E e) {
        int level = dummyHead.next[0] == null ? 1 : randomLevel();
        Node<E>[] prevNodes = (Node<E>[]) new Node[level];
        for (int i = 0; i < level; i++) {
            prevNodes[i] = dummyHead;
        }

        Node<E> prev = dummyHead;
        for (int i = maxLevel - 1; i >= 0; i--) {
            while (prev.next[i] != null && prev.next[i].e.compareTo(e) < 0) {
                prev = prev.next[i];
            }

            if (i < level) {
                prevNodes[i] = prev;
            }
        }

        if (prev.next[0] != null && prev.next[0].e.compareTo(e) == 0) {
            return;
        }

        Node<E> node = new Node<>(e);
        for (int i = 0; i < level; i++) {
            node.next[i] = prevNodes[i].next[i];
            prevNodes[i].next[i] = node;
        }

        if (maxLevel < level) {
            maxLevel = level;
        }
    }

    private int randomLevel() {
        int level = 1;
        for (int i = 1; i < MAX_LEVEL; i++) {
            if (random.nextInt() % 2 == 1) {
                level++;
            }
        }

        return level;
    }

    @SuppressWarnings("all")
    public void remove(E e) {
        Node<E>[] prevNodes = (Node<E>[]) new Node[maxLevel];
        Node<E> prev = dummyHead;
        for (int i = maxLevel - 1; i >= 0; i--) {
            while (prev.next[i] != null && prev.next[i].e.compareTo(e) < 0) {
                prev = prev.next[i];
            }

            prevNodes[i] = prev;
        }

        if (prev.next[0] != null && prev.next[0].e.compareTo(e) == 0) {
            for (int i = maxLevel - 1; i >= 0; i--) {
                Node<E> delNode = prevNodes[i].next[i];
                if (delNode != null && delNode.e.compareTo(e) == 0) {
                    prevNodes[i].next[i] = delNode.next[i];
                    delNode.next[i] = null;
                }
            }
        }
    }
}
